# 该文件生成21张效果图，用于比较模型的训练效果
# 与_test不同的是,它的G达到了7层

from __future__ import print_function

import warnings
warnings.filterwarnings("ignore")
import glob
import os
from os import sep
from os.path import join, exists, splitext
import time
import tensorflow as tf
import numpy as np
from scipy.misc import imread, imsave
from Generator import Generator
import  argparse

# 定义一下文件夹名



def generate(ir_path, vis_path, model_path, name, output_path):
	ir_img = imread(ir_path) / 255.0
	vis_img = imread(vis_path) / 255.0
	ir_dimension = list(ir_img.shape)
	vis_dimension = list(vis_img.shape)
	ir_dimension.insert(0, 1)
	ir_dimension.append(1)			# [1, 268, 360, 1]
	vis_dimension.insert(0, 1)
	vis_dimension.append(1)			# [1, 268, 360, 1]
	ir_img = ir_img.reshape(ir_dimension)		# 将(268, 360)转为(1, 268, 360, 1)
	vis_img = vis_img.reshape(vis_dimension)

	os.environ["CUDA_VISIBLE_DEVICES"] = '1'
	# 按照固定的比例分配。
	config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
	# 以下代码会占用所有可使用的GPU的50%显存，已试过80%
	config.gpu_options.per_process_gpu_memory_fraction = 0.5
	config.gpu_options.allow_growth = True

	with tf.Graph().as_default(), tf.Session(config=config) as sess:
		SOURCE_VIS = tf.placeholder(tf.float32, shape = vis_dimension, name = 'SOURCE_VIS')
		SOURCE_ir = tf.placeholder(tf.float32, shape = ir_dimension, name = 'SOURCE_ir')

		G = Generator('Generator')
		output_image = G.transform(vis = SOURCE_VIS, ir = SOURCE_ir)


		# restore the trained model and run the style transferring
		saver = tf.train.Saver()
		saver.restore(sess, model_path)
		output = sess.run(output_image, feed_dict = {SOURCE_VIS: vis_img, SOURCE_ir: ir_img})
		output = output[0, :, :, 0]
		imsave(output_path + name + '.png', output)



def save_images(paths, datas, save_path, prefix = None, suffix = None):
	if isinstance(paths, str):
		paths = [paths]

	assert (len(paths) == len(datas))

	if not exists(save_path):
		mkdir(save_path)

	if prefix is None:
		prefix = ''
	if suffix is None:
		suffix = ''

	for i, path in enumerate(paths):
		data = datas[i]
		# print('data ==>>\n', data)
		if data.shape[2] == 1:
			data = data.reshape([data.shape[0], data.shape[1]])
		# print('data reshape==>>\n', data)

		name, ext = splitext(path)
		name = name.split(sep)[-1]

		path = join(save_path, prefix + suffix + ext)
		print('data path==>>', path)
		imsave(path, data)

def mkdir(dir):
	"这个函数创建不存在的 路径文件夹"
	if not os.path.exists(dir):
		os.makedirs(dir)

def getRealImgs(data_dir):	 
	data = glob.glob(os.path.join(data_dir, "*.png"))
	data.extend(glob.glob(os.path.join(data_dir, "*.jpg")))
	data.sort(key=lambda x: x[len(data_dir):-4])
	return data

def gen_img(model_path,irimgs_path,visimgs_path,output_path):
	print('\nBegin to generate pictures ...\n')
	mkdir(output_path)
	Time=[]

	for ir_path,vis_path in zip(getRealImgs(irimgs_path),getRealImgs(visimgs_path)):
		name = ir_path[len(irimgs_path):-4]
		begin = time.time()

		generate(ir_path, vis_path, model_path, name, output_path)

		end = time.time()
		Time.append(end - begin)
		print(f"	pic_num: {name} 生成完毕")
	print("Time: mean:%s, std: %s" % (np.mean(Time), np.std(Time)))


parser = argparse.ArgumentParser()
parser.add_argument("-m","--model",help="introduce model's path")
parser.add_argument("-ir","--ir_path",help="intrduce path")
parser.add_argument("-vi","--vi_path",help="introduce  visible imgs' path")
parser.add_argument("-p","--result_path",help="introduce  fused imgs' path")

if __name__ == '__main__':

	args = parser.parse_args()
	MODEL_PATH = args.model
	ir_path = args.ir_path
	vi_path= args.vi_path
	RESULT_PATH = args.result_path

	if not (MODEL_PATH and ir_path and vi_path and RESULT_PATH):
		exit(-1)
	# 调用MODEL_PATH处的模型，从path处读取图像，融合后，写入results文件夹中
	try:
		gen_img(MODEL_PATH, ir_path, vi_path, RESULT_PATH)
	except Exception as err:
		print(err)
